# -*- coding: utf-8 -*-
# @Author  : gaoyu
# @Time    : 2025/9/4 17:06
# @Function:
import pandas as pd
import geopandas as gpd
from loguru import logger
from sqlalchemy import create_engine, inspect, text


def clear_database_tables(db_url, tables):
    """
    清空指定的数据库表。

    Args:
        db_url (str): 数据库连接字符串。
        tables (list): 要清空的表名列表。
    """
    engine = create_engine(db_url)
    logger.info("--- 正在清空旧数据表 ---")
    with engine.connect() as conn:
        for table_name in tables:
            logger.info(f"正在删除表 '{table_name}'...")
            conn.execute(text(f"DROP TABLE IF EXISTS {table_name} CASCADE;"))
            conn.commit()
            logger.info(f"表 '{table_name}' 删除成功。")


def create_table_if_not_exists(engine, table_name, gdf, primary_key=None):
    """
    检查PostGIS数据库中是否存在指定的表。如果不存在，则创建新表，
    为几何字段添加GIST索引，并可选择添加主键约束。

    Args:
        engine (sqlalchemy.engine.Engine): 数据库连接引擎。
        table_name (str): 要创建或检查的表名。
        gdf (geopandas.GeoDataFrame): 用于定义表结构的GeoDataFrame。
        primary_key (str, optional): 指定用作主键的列名。默认为None。
    """
    with engine.connect() as conn:
        inspector = inspect(engine)
        if not inspector.has_table(table_name):
            logger.info(f"表 '{table_name}' 不存在，正在创建...")
            # to_postgis 可以自动创建表，并移除 'srid' 参数
            gdf.to_postgis(table_name, conn, if_exists='replace', index=False)

            # 添加 GIST 索引，用于优化空间查询
            logger.info(f"正在为表 '{table_name}' 添加 GIST 索引...")
            conn.execute(
                text(f"CREATE INDEX {table_name}_geom_idx ON {table_name} USING GIST (geometry);")
            )
            conn.commit()
            logger.info(f"GIST 索引添加成功。")

            if primary_key:
                # 添加主键约束，确保数据的唯一性
                logger.info(f"正在为表 '{table_name}' 添加主键约束 '{primary_key}'...")
                try:
                    conn.execute(
                        text(f"ALTER TABLE {table_name} ADD PRIMARY KEY ({primary_key});")
                    )
                    conn.commit()
                    logger.info(f"主键约束添加成功。")
                except Exception as e:
                    logger.warning(f"添加主键约束失败: {e}")
                    logger.info("可能是因为数据中存在重复值，将跳过主键约束的添加。")
        else:
            logger.info(f"表 '{table_name}' 已存在。")
            # 检查主键是否存在
            if primary_key:
                try:
                    result = conn.execute(
                        text(f"SELECT constraint_name FROM information_schema.table_constraints WHERE table_name = '{table_name}' AND constraint_type = 'PRIMARY KEY';")
                    )
                    pk_exists = result.fetchone() is not None
                    if not pk_exists:
                        logger.info(f"表 '{table_name}' 缺少主键约束，尝试添加...")
                        try:
                            conn.execute(
                                text(f"ALTER TABLE {table_name} ADD PRIMARY KEY ({primary_key});")
                            )
                            conn.commit()
                            logger.info(f"主键约束添加成功。")
                        except Exception as e:
                            logger.warning(f"添加主键约束失败: {e}")
                            logger.info("可能是因为数据中存在重复值，将跳过主键约束的添加。")
                except Exception as e:
                    logger.warning(f"检查主键约束时出错: {e}")


def import_provinces_data(db_url, province_files, target_srid=4326):
    """
    导入省份数据到 PostGIS 数据库。该函数会自动处理多个文件，
    统一几何类型并导入到同一张表中。每个省份zip文件对应数据库中的一行。

    Args:
        db_url (str): 数据库连接字符串。
        province_files (list): 包含所有省份数据文件路径的列表。
        target_srid (int, optional): 目标空间参考系统ID。默认为4326 (WGS84)。
    """
    engine = create_engine(db_url)
    province_table_name = 'gis_provinces'
    logger.info(f"--- 导入省份数据到表 '{province_table_name}' ---")

    all_provinces = []
    
    for file in province_files:
        logger.info(f"正在读取文件: '{file}'")
        try:
            # 尝试以UTF-8编码读取，以解决中文乱码问题
            gdf = gpd.read_file(file, encoding='utf-8')
            gdf.columns = gdf.columns.str.lower()
            logger.info(f"文件 '{file}' 包含 {len(gdf)} 个几何对象")
        except Exception as e:
            logger.error(f"无法读取文件 '{file}'，请检查文件路径和编码。错误信息: {e}")
            continue

        # 检查并转换坐标系（如果需要），确保数据统一
        if gdf.crs and gdf.crs.to_epsg() != target_srid:
            logger.info(f"文件 '{file}' 的坐标系是 {gdf.crs.to_epsg()}，正在转换为 EPSG:{target_srid}。")
            gdf = gdf.to_crs(epsg=target_srid)

        # 确保几何类型为 MultiPolygon，但不使用explode()来保持每个省份为一行
        from shapely.geometry import MultiPolygon, Polygon
        
        def ensure_multipolygon(geom):
            if isinstance(geom, Polygon):
                return MultiPolygon([geom])
            elif isinstance(geom, MultiPolygon):
                return geom
            else:
                logger.warning(f"未知的几何类型: {type(geom)}，尝试转换为MultiPolygon")
                return MultiPolygon([geom]) if hasattr(geom, 'geoms') else geom
        
        gdf['geometry'] = gdf['geometry'].apply(ensure_multipolygon)
        logger.info(f"几何类型统一为MultiPolygon完成，共 {len(gdf)} 行数据")
        
        # 检查gb字段是否存在
        if 'gb' not in gdf.columns:
            logger.warning(f"文件 '{file}' 中未找到 'gb' 字段，将使用其他可用字段作为标识")
            # 如果没有gb字段，尝试使用其他可能的标识字段
            possible_id_fields = ['code', 'id', 'name', 'province_code']
            for field in possible_id_fields:
                if field in gdf.columns:
                    gdf['gb'] = gdf[field]
                    logger.info(f"使用字段 '{field}' 作为gb字段")
                    break
            else:
                # 如果都没有，创建一个基于文件名的标识
                import os
                file_name = os.path.splitext(os.path.basename(file))[0]
                gdf['gb'] = file_name.replace('_省', '')
                logger.info(f"创建gb字段，值为: {gdf['gb'].iloc[0]}")
        
        # 检查gb字段的唯一性
        if gdf['gb'].duplicated().any():
            logger.warning(f"文件 '{file}' 中gb字段存在重复值，将保留第一个")
            gdf = gdf.drop_duplicates(subset=['gb'], keep='first')
        
        all_provinces.append(gdf)
    
    if not all_provinces:
        logger.error("没有成功读取任何省份数据文件")
        return
    
    # 合并所有省份数据
    combined_gdf = gpd.GeoDataFrame(pd.concat(all_provinces, ignore_index=True))
    logger.info(f"合并后共有 {len(combined_gdf)} 个省份")
    
    # 再次检查gb字段的全局唯一性
    if combined_gdf['gb'].duplicated().any():
        logger.warning("合并后的数据中gb字段存在重复值，将删除重复项")
        duplicates = combined_gdf[combined_gdf['gb'].duplicated(keep=False)]
        logger.info(f"重复的gb值: {duplicates['gb'].tolist()}")
        combined_gdf = combined_gdf.drop_duplicates(subset=['gb'], keep='first')
        logger.info(f"删除重复项后剩余 {len(combined_gdf)} 个省份")
    
    # 创建表并设置gb为主键
    create_table_if_not_exists(engine, province_table_name, combined_gdf, primary_key='gb')
    
    # 导入数据
    try:
        combined_gdf.to_postgis(province_table_name, engine, if_exists='replace', index=False)
        logger.info(f"成功导入 {len(combined_gdf)} 个省份到表 '{province_table_name}'")
    except Exception as e:
        logger.error(f"导入省份数据失败: {e}")
        # 尝试使用replace模式
        try:
            logger.info("尝试使用replace模式重新导入...")
            combined_gdf.to_postgis(province_table_name, engine, if_exists='replace', index=False)
            # 重新添加索引和主键
            with engine.connect() as conn:
                conn.execute(text(f"CREATE INDEX IF NOT EXISTS {province_table_name}_geom_idx ON {province_table_name} USING GIST (geometry);"))
                conn.execute(text(f"ALTER TABLE {province_table_name} ADD PRIMARY KEY (gb);"))
                conn.commit()
            logger.info(f"使用replace模式成功导入 {len(combined_gdf)} 个省份")
        except Exception as e2:
            logger.error(f"使用replace模式导入也失败: {e2}")

        logger.info(f"'{file}' 数据导入完成。")

    logger.info("所有省份数据导入完成。")


def import_reservoirs_data(db_url, reservoir_file, target_srid=4326):
    """
    导入水库数据到 PostGIS 数据库。该函数会处理单个文件，
    统一几何类型并为表添加主键约束。

    Args:
        db_url (str): 数据库连接字符串。
        reservoir_file (str): 水库数据文件的路径。
        target_srid (int, optional): 目标空间参考系统ID。默认为4326 (WGS84)。
    """
    engine = create_engine(db_url)
    reservoir_table_name = 'gis_reservoirs'
    logger.info(f"\n--- 导入水库数据到表 '{reservoir_table_name}' ---")

    try:
        # 尝试默认编码读取
        gdf_reservoirs = gpd.read_file(reservoir_file)
        gdf_reservoirs.columns = gdf_reservoirs.columns.str.lower()
    except Exception as e:
        logger.error(f"无法读取水库文件 '{reservoir_file}'，请检查文件路径和编码。错误信息: {e}")
        return

    # 检查并转换坐标系（如果需要），确保数据统一
    if gdf_reservoirs.crs and gdf_reservoirs.crs.to_epsg() != target_srid:
        logger.info(f"水库文件坐标系是 {gdf_reservoirs.crs.to_epsg()}，正在转换为 EPSG:{target_srid}。")
        gdf_reservoirs = gdf_reservoirs.to_crs(epsg=target_srid)

    # 首先删除原始数据中的重复stcd，保留第一个出现的记录
    logger.info(f"原始数据包含 {len(gdf_reservoirs)} 条记录")
    initial_duplicates = gdf_reservoirs.duplicated(subset=['stcd']).sum()
    if initial_duplicates > 0:
        logger.info(f"发现 {initial_duplicates} 条重复的stcd记录，正在删除...")
        gdf_reservoirs = gdf_reservoirs.drop_duplicates(subset=['stcd'], keep='first')
        logger.info(f"删除重复后剩余 {len(gdf_reservoirs)} 条记录")

    # 统一几何类型为 MultiPolygon，以兼容Polygon和MultiPolygon类型
    # 注意：explode 会为 MultiPolygon 的每个部分创建一个新行，
    # 但所有新行都会继承原始行的属性，包括 stcd。
    gdf_reservoirs = gdf_reservoirs.explode(index_parts=False).reset_index(drop=True)
    logger.info(f"explode操作后包含 {len(gdf_reservoirs)} 条记录")

    # 在explode后再次删除重复的行，以避免主键冲突
    # 必须在 explode 之后执行此操作，因为 explode 可能会引入重复的 stcd
    explode_duplicates = gdf_reservoirs.duplicated(subset=['stcd']).sum()
    if explode_duplicates > 0:
        logger.info(f"explode后发现 {explode_duplicates} 条重复的stcd记录，正在删除...")
        gdf_reservoirs = gdf_reservoirs.drop_duplicates(subset=['stcd'], keep='first')
        logger.info(f"最终处理后剩余 {len(gdf_reservoirs)} 条记录")

    # 检查是否还有空的stcd值
    null_stcd_count = gdf_reservoirs['stcd'].isnull().sum()
    if null_stcd_count > 0:
        logger.warning(f"发现 {null_stcd_count} 条记录的stcd为空，将被过滤掉")
        gdf_reservoirs = gdf_reservoirs.dropna(subset=['stcd'])
        logger.info(f"过滤空stcd后剩余 {len(gdf_reservoirs)} 条记录")

    # 检查表是否存在，并创建（如果需要），同时添加 GIST 索引和主键
    create_table_if_not_exists(engine, reservoir_table_name, gdf_reservoirs, primary_key='stcd')

    # 导入水库数据
    try:
        gdf_reservoirs.to_postgis(reservoir_table_name, engine, if_exists='replace', index=False)
        logger.info(f"'{reservoir_file}' 数据导入完成，成功导入 {len(gdf_reservoirs)} 条记录。")
    except Exception as e:
        logger.error(f"导入数据时发生错误: {e}")
        # 如果仍然有主键冲突，尝试使用replace模式
        logger.info("尝试使用replace模式重新导入...")
        try:
            gdf_reservoirs.to_postgis(reservoir_table_name, engine, if_exists='replace', index=False)
            logger.info(f"使用replace模式成功导入 {len(gdf_reservoirs)} 条记录。")
        except Exception as e2:
            logger.error(f"replace模式导入也失败: {e2}")
            return

    logger.info("\n所有水库数据导入任务完成！")


def calculate_attributes_and_print(db_url, table_name, display_cols, primary_key='stcd'):
    """
    计算表中所有要素的面积、周长和中心点，并输出统计摘要。
    此函数直接从PostGIS读取数据，并在Python中进行计算和统计。

    Args:
        db_url (str): 数据库连接字符串。
        table_name (str): 要处理的表名。
        display_cols (list): 要打印的显示列名列表。
        primary_key (str, optional): 主键列名，用于唯一标识每行。默认为'stcd'。
    """
    engine = create_engine(db_url)
    logger.info(f"--- 正在为表 '{table_name}' 计算几何属性 ---")

    try:
        # 从数据库读取数据
        gdf = gpd.read_postgis(f"SELECT * FROM {table_name}", engine, geom_col='geometry')
        if gdf.empty:
            logger.warning(f"警告：表 '{table_name}' 为空，无法计算属性。")
            return

        # 将坐标系转换为投影坐标系以进行准确的面积和周长计算（单位为平方米）
        projected_gdf = gdf.to_crs(gdf.estimate_utm_crs())

        # 计算面积、周长和中心点
        gdf['area_m2'] = projected_gdf.area
        gdf['perimeter_m'] = projected_gdf.length

        # 计算中心点并转换回 WGS84 坐标系
        centroids = projected_gdf.centroid.to_crs(epsg=4326)
        gdf['centroid_x'] = centroids.x
        gdf['centroid_y'] = centroids.y

        # 输出统计摘要
        total_count = len(gdf)
        total_area = gdf['area_m2'].sum()
        avg_area = gdf['area_m2'].mean()
        max_area = gdf['area_m2'].max()
        min_area = gdf['area_m2'].min()
        
        logger.info(f"✅ 计算完成 - 共处理 {total_count} 个要素")
        logger.info(f"📊 面积统计: 总计 {total_area/1e6:.2f} km², 平均 {avg_area/1e6:.2f} km², 最大 {max_area/1e6:.2f} km², 最小 {min_area/1e6:.2f} km²")
        
        # 只显示前3个和后3个要素的详细信息
        if total_count > 6:
            logger.info("📋 详细信息 (显示前3个和后3个):")
            sample_indices = list(range(3)) + list(range(total_count-3, total_count))
        else:
            logger.info("📋 详细信息:")
            sample_indices = list(range(total_count))
            
        for i, (index, row) in enumerate(gdf.iterrows()):
            if i in sample_indices:
                display_info = ', '.join([f"{col}: {row[col]}" for col in display_cols])
                logger.info(f"   {row[primary_key]} | {display_info} | {row['area_m2']/1e6:.2f} km² | ({row['centroid_x']:.4f}, {row['centroid_y']:.4f})")
            elif i == 3 and total_count > 6:
                logger.info(f"   ... (省略中间 {total_count-6} 个要素) ...")

    except Exception as e:
        logger.error(f"计算几何属性时发生错误: {e}")


def find_reservoirs_in_provinces(db_url, province_table, reservoir_table):
    """
    执行空间连接，找出每个省份所包含的水库，并输出统计摘要。

    Args:
        db_url (str): 数据库连接字符串。
        province_table (str): 省份表名。
        reservoir_table (str): 水库表名。
    """
    engine = create_engine(db_url)
    logger.info(f"\n--- 空间查询：省份包含水库关系分析 ---")

    try:
        # 统计查询
        stats_query = text(f"""
            SELECT 
                prov.name AS province_name,
                COUNT(res.stcd) AS reservoir_count
            FROM 
                {province_table} AS prov
            LEFT JOIN {reservoir_table} AS res ON ST_Within(res.geometry, prov.geometry)
            GROUP BY prov.name, prov.gb
            ORDER BY reservoir_count DESC;
        """)
        
        # 详细查询（限制前几个结果）
        detail_query = text(f"""
            SELECT 
                prov.name AS province_name,
                res.stnm AS reservoir_name,
                res.stcd AS stcd
            FROM 
                {reservoir_table} AS res,
                {province_table} AS prov
            WHERE 
                ST_Within(res.geometry, prov.geometry)
            LIMIT 10;
        """)

        with engine.connect() as conn:
            # 获取统计结果
            stats_result = conn.execute(stats_query).fetchall()
            detail_result = conn.execute(detail_query).fetchall()
            
            if not detail_result:
                logger.info("❌ 未找到任何包含关系")
                return

            # 输出统计摘要
            total_relations = len(detail_result)
            provinces_with_reservoirs = len([r for r in stats_result if r.reservoir_count > 0])
            
            logger.info(f"✅ 包含关系统计: 共 {total_relations} 个水库被省份包含")
            logger.info(f"📊 省份统计: {provinces_with_reservoirs} 个省份包含水库")
            
            # 显示各省份的水库数量
            logger.info("📋 各省份包含的水库数量:")
            for row in stats_result:
                if row.reservoir_count > 0:
                    logger.info(f"   {row.province_name}: {row.reservoir_count} 个水库")
            
            # 显示前几个具体的包含关系
            if total_relations > 0:
                logger.info(f"📝 具体包含关系 (显示前{min(10, total_relations)}个):")
                for i, row in enumerate(detail_result[:10]):
                    logger.info(f"   {i+1}. {row.province_name} ← {row.reservoir_name} ({row.stcd})")
                if total_relations > 10:
                    logger.info(f"   ... (还有 {total_relations-10} 个包含关系) ...")

    except Exception as e:
        logger.error(f"执行空间查询时发生错误: {e}")


def find_intersecting_reservoirs_and_provinces(db_url, province_table, reservoir_table):
    """
    执行空间连接，找出与每个省份相交的水库，并输出统计摘要。

    Args:
        db_url (str): 数据库连接字符串。
        province_table (str): 省份表名。
        reservoir_table (str): 水库表名。
    """
    engine = create_engine(db_url)
    logger.info(f"\n--- 空间查询：省份与水库相交关系分析 ---")

    try:
        # 统计查询
        stats_query = text(f"""
            SELECT 
                prov.name AS province_name,
                COUNT(res.stcd) AS reservoir_count
            FROM 
                {province_table} AS prov
            LEFT JOIN {reservoir_table} AS res ON ST_Intersects(res.geometry, prov.geometry)
            GROUP BY prov.name, prov.gb
            ORDER BY reservoir_count DESC;
        """)
        
        # 详细查询（限制前几个结果）
        detail_query = text(f"""
            SELECT 
                prov.name AS province_name,
                res.stnm AS reservoir_name,
                res.stcd AS stcd
            FROM 
                {reservoir_table} AS res,
                {province_table} AS prov
            WHERE 
                ST_Intersects(res.geometry, prov.geometry)
            LIMIT 10;
        """)

        with engine.connect() as conn:
            # 获取统计结果
            stats_result = conn.execute(stats_query).fetchall()
            detail_result = conn.execute(detail_query).fetchall()
            
            if not detail_result:
                logger.info("❌ 未找到任何相交关系")
                return

            # 输出统计摘要
            total_relations = len(detail_result)
            provinces_with_reservoirs = len([r for r in stats_result if r.reservoir_count > 0])
            
            logger.info(f"✅ 相交关系统计: 共 {total_relations} 个水库与省份相交")
            logger.info(f"📊 省份统计: {provinces_with_reservoirs} 个省份与水库相交")
            
            # 显示各省份的相交水库数量
            logger.info("📋 各省份相交的水库数量:")
            for row in stats_result:
                if row.reservoir_count > 0:
                    logger.info(f"   {row.province_name}: {row.reservoir_count} 个水库")
            
            # 显示前几个具体的相交关系
            if total_relations > 0:
                logger.info(f"📝 具体相交关系 (显示前{min(10, total_relations)}个):")
                for i, row in enumerate(detail_result[:10]):
                    logger.info(f"   {i+1}. {row.province_name} ∩ {row.reservoir_name} ({row.stcd})")
                if total_relations > 10:
                    logger.info(f"   ... (还有 {total_relations-10} 个相交关系) ...")

    except Exception as e:
        logger.error(f"执行空间查询时发生错误: {e}")


if __name__ == '__main__':
    # 请将以下连接信息替换为你的数据库配置
    DB_URL = "postgresql://postgres:water@10.48.0.85:5432/hydropulse"

    # SHP 文件列表 - 从项目根目录读取
    import os
    project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
    province_files = [
        os.path.join(project_root, 'data', '吉林省_省.zip'),
        os.path.join(project_root, 'data', '辽宁省_省.zip'),
        os.path.join(project_root, 'data', '黑龙江省_省.zip')
    ]
    reservoir_file = os.path.join(project_root, 'data', 'merged_output.shp')

    # --- 0. 清空数据库表 ---
    clear_database_tables(DB_URL, ['gis_provinces', 'gis_reservoirs'])

    # --- 1. 导入原始数据到数据库 ---
    import_provinces_data(DB_URL, province_files, target_srid=4326)
    import_reservoirs_data(DB_URL, reservoir_file, target_srid=4326)

    # --- 2. 计算几何属性并打印 ---
    calculate_attributes_and_print(DB_URL, 'gis_provinces', display_cols=['name', 'gb'], primary_key='gb')
    calculate_attributes_and_print(DB_URL, 'gis_reservoirs', display_cols=['stnm'], primary_key='stcd')

    # --- 3. 执行空间连接并打印结果 ---
    find_reservoirs_in_provinces(DB_URL, 'gis_provinces', 'gis_reservoirs')
    find_intersecting_reservoirs_and_provinces(DB_URL, 'gis_provinces', 'gis_reservoirs')
